package mlp;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.PrintStream;
import java.util.BitSet;

import javax.swing.JOptionPane;

import digitRecognitionProblem.DigitRecognitionMlpFitness;

public class TrainTrimmedSparseMlp {

	private static int trainSize  = 38000;
	
	// indices for main arguments
	private enum ArgNums {
		TEST_FILE, TRAIN_FILE, START_DIGIT, END_DIGIT, LEARNING_RATE, MOMENTUM, ERROR_RATE, TRAIN_SIZE, FROM_INDEX, EDGE_SWITCH_DIR, TRAINED_DIR
	}
	
	public static void main(String[] args) {
		try {
			File file = new File("Details.txt");
			FileOutputStream fos;
			fos = new FileOutputStream(file);
			PrintStream ps = new PrintStream(fos);
			System.setOut(ps);
			System.setErr(ps);
		} catch (FileNotFoundException e1) {
			e1.printStackTrace();
		}
		
		try {
			
			String testFileName = args[ArgNums.TEST_FILE.ordinal()];
			String trainFileName = args[ArgNums.TRAIN_FILE.ordinal()];
			int startDigit = Integer.parseInt(args[ArgNums.START_DIGIT.ordinal()]);
			int endDigit = Integer.parseInt(args[ArgNums.END_DIGIT.ordinal()]);
			Mlp.LEARNING_RATE = Float.parseFloat(args[ArgNums.LEARNING_RATE.ordinal()]);
			Mlp.MOMENTUM = Float.parseFloat(args[ArgNums.MOMENTUM.ordinal()]);
			Mlp.ERROR_RATE = Float.parseFloat(args[ArgNums.ERROR_RATE.ordinal()]);
			trainSize = Integer.parseInt(args[ArgNums.TRAIN_SIZE.ordinal()]);
			int fromIndex = Integer.parseInt(args[ArgNums.FROM_INDEX.ordinal()]);
			String edgeSwitchDir = args[ArgNums.EDGE_SWITCH_DIR.ordinal()];
			String trainedDir = args[ArgNums.TRAINED_DIR.ordinal()];
			Mlp.NUMBER_EXAMPLE = trainSize;

			Mlp.NUMBER_EXAMPLE = trainSize;
			Mlp.INIT_WEIGHT_RANGE = 0.125f;
			
			// document arguments
			System.out.println("Test file = " + testFileName);
			System.out.println("Train file = " + trainFileName);
			System.out.println("Starts from digit = " + startDigit);
			System.out.println("Ends at digit = " + endDigit);
			System.out.println("Learning rate = " + Mlp.LEARNING_RATE);
			System.out.println("Momentum = " + Mlp.MOMENTUM);
			System.out.println("Error rate = " + Mlp.ERROR_RATE);
			System.out.println("Train size = " + trainSize);
			System.out.println("From index = " + fromIndex);
			System.out.println("From index = " + fromIndex);
			System.out.println("From index = " + fromIndex);
			System.out.println("******************************************");
			
			for (int digit = startDigit ; digit <= endDigit ; ++digit) {

				DigitRecognitionMlpFitness.initTestData(testFileName, digit);

				// create directory for network learning current digit
				String path = "" + digit;
				File dir = new File(path);
				dir.mkdir();
				
				DigitRecognitionMlpFitness.initTestData(testFileName, digit);
				System.out.println("************* MLP recognizing "+ digit +" *************");
				MlpExamples exam = new MlpExamples(trainFileName, fromIndex*trainSize, fromIndex*trainSize+trainSize, digit);
				
				// read trained network
				SparseMlp mlp = null;
				try {
					FileInputStream readFile = new FileInputStream(new File(trainedDir + "/" + path + "/mlp_0_rec_" + digit + ".txt"));
					ObjectInputStream fileStream = new ObjectInputStream(readFile);
					mlp = (SparseMlp)fileStream.readObject();
					fileStream.close();
				} catch (Exception e) {
					e.printStackTrace();
				}

				// build path to edge switch
				String edgeSwithFile = edgeSwitchDir + "/edge_switch_rec_" + digit + ".txt";
				
				// trim network
				trimEdges(edgeSwithFile, mlp);
				
				// re-train the network
				mlp.learn(exam.getInput(), exam.getOutput());
				
				// test the network
				DigitRecognitionMlpFitness funcFit = new DigitRecognitionMlpFitness(mlp);
						
				System.out.println("MLP (DIGIT "+digit+") :  Success rate on test set: "+ funcFit.getFitness(null));				
				System.out.println("\n************************************");
				FileOutputStream writeFile;
				try {
					writeFile = new FileOutputStream(path + "/" + "trimmed_mlp_rec_" + digit +".txt");
					ObjectOutputStream fileStream = new ObjectOutputStream(writeFile);
					fileStream.writeObject(mlp);
					fileStream.close();
				} catch (FileNotFoundException e) {
					e.printStackTrace();
				} catch (IOException e) {
					e.printStackTrace();
				}
			}			

			JOptionPane.showMessageDialog(null, "Building the networks is finished");
		} 
		catch (Exception e) {
			System.out.println("***********************ERROR*********************************");
			System.out.println(e.getMessage()+"\n");
			System.out.println("**************************************************************");
			JOptionPane.showMessageDialog(null, "ERROR, see Details.txt", "Error", JOptionPane.ERROR_MESSAGE);
		}
	}
	
	
	private static void trimEdges(String edgeSwithFile, SparseMlp mlp) {
		
		// read serialized edge switch
		BitSet edgeSwitch = null;
		try {
			FileInputStream readFile = new FileInputStream(new File(edgeSwithFile));
			ObjectInputStream fileStream = new ObjectInputStream(readFile);
			edgeSwitch = (BitSet)fileStream.readObject();
			fileStream.close();
		} catch (Exception e) {
			e.printStackTrace();
		} 

		// trim edges according to parsed data
		mlp.setSwitch(edgeSwitch);
	}
}
